import networkx as nx
import matplotlib.pyplot as plt
import time
from collections import deque
import numpy as np
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

def generate_gas_network():
    G = nx.DiGraph()
    nodes = []
    for i in range(1, 3+1):
        nodes.append((i, {'type': 'source'}))   # 气源节点
    for i in range(4, 11+1):
        nodes.append((i, {'type': 'other'}))    # 其他节点（调压站等）    
    for i in range(12, 18+1):
        nodes.append((i, {'type': 'valve'}))    # 阀门节点
    for i in range(40, 41+1):
        nodes.append((i, {'type': 'valve'}))    # 阀门节点
    for i in range(30, 30+1):
        nodes.append((i, {'type': 'leak', 'leak_weight': 1}))    # 泄漏节点
    for i in range(32, 33+1):
        nodes.append((i, {'type': 'leak', 'leak_weight': 1}))    # 泄漏节点        
    for i in range(20, 29+1):
        nodes.append((i, {'type': 'user', 'user_weight': 1}))    # 用户节点（调压箱）
    for i in range(50, 50+1):
        nodes.append((i, {'type': 'user', 'user_weight': 1}))    # 用户节点（调压箱）      
    for i in range(31, 31+1):
        nodes.append((i, {'type': 'user', 'user_weight': 1}))    # 用户节点（调压箱）                
    pos = {
        1: np.array([-0.2, 0.2]),
        2: np.array([0.2, 0.2]),
        3: np.array([0, -0.2]),
        4: np.array([-0.4, 0.1]),
        5: np.array([-0.3, -0.22]),
        6: np.array([0.3, -0.22]),
        7: np.array([0.4, 0.1]),
        8: np.array([0.1, 0.1]),
        9: np.array([-0.1, 0.1]),
        33: np.array([-0.02, 0.1]),
        10: np.array([-0.2, -0.1]),
        11: np.array([0.2, -0.1]),
        12: np.array([-0.4, -0.1]),
        13: np.array([0.4, -0.1]),
        14: np.array([0, -0.1]),
        15: np.array([0.2, 0.1]),
        16: np.array([-0.1, 0]),
        17: np.array([-0.3, 0]),
        18: np.array([0.25, -0.15]),
        40: np.array([0.3, 0.15]),
        41: np.array([0.4, 0.05]),
        30: np.array([0.2, 0]),
        31: np.array([0.3, -0.15]),
        32: np.array([0.07, 0.05]),
        20: np.array([0.3, 0]),
        21: np.array([0.3, -0.05]),
        22: np.array([0, -0.05]),
        23: np.array([0, 0]),
        24: np.array([0, 0.05]),
        25: np.array([0.07, 0]),
        26: np.array([0.07, -0.05]),
        27: np.array([-0.3, 0.1]),
        28: np.array([-0.2, 0.1]),
        29: np.array([0.4, -0.22]),
        50: np.array([0.4, -0.05])
    }
    G.add_edge(1, 4)
    G.add_edge(4, 12)
    G.add_edge(12, 5)
    G.add_edge(4, 17)
    G.add_edge(10, 16)
    G.add_edge(5, 10)
    G.add_edge(3, 5)
    G.add_edge(13, 6)
    G.add_edge(7, 41)
    G.add_edge(41, 50)
    G.add_edge(50, 13)
    G.add_edge(6, 18)
    G.add_edge(6, 29)
    G.add_edge(18, 11)
    G.add_edge(10, 14)
    G.add_edge(11, 8)
    G.add_edge(7, 15)
    G.add_edge(15, 8)
    G.add_edge(8, 33)
    G.add_edge(33, 9)
    G.add_edge(2, 40)
    G.add_edge(40, 7)
    G.add_edge(11, 30)
    G.add_edge(30, 20)
    G.add_edge(20, 21)
    G.add_edge(21, 31)
    G.add_edge(6, 31)
    G.add_edge(14, 22)
    G.add_edge(22, 23)
    G.add_edge(23, 24)
    G.add_edge(24, 32)
    G.add_edge(32, 25)
    G.add_edge(25, 26)
    G.add_edge(26, 22)
    G.add_edge(25, 8)
    G.add_edge(17, 27)
    G.add_edge(27, 28)
    G.add_edge(28, 9)
    G.add_edge(16, 24)
    G.add_nodes_from(nodes)
    return G, pos

def get_outward_connected_components(G, nodes):
    """
    从给定的节点集合 nodes 出发，通过出边遍历，返回可达的节点集合和边集合。
    
    参数:
        G (nx.DiGraph): 有向图
        nodes (list): 初始节点集合
    
    返回:
        reachable_nodes (set): 可达的节点集合
        no_reachable_edges (set): 可达的边集合（格式为元组 (u, v)）
    """
    # 1. 找到所有可达的节点
    visited = set()
    queue = deque(nodes)
    
    for node in nodes:
        if node not in visited:
            visited.add(node)
    
    while queue:
        current = queue.popleft()
        for neighbor in G.successors(current):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)
    
    # 2. 收集所有可达的边（起点和终点都在 reachable_nodes 中的边）
    reachable_edges = set()
    for u, v in G.edges:
        if u in visited and v in visited:
            reachable_edges.add((u, v))
    
    return visited, reachable_edges

def visualize_network(H, pos, name, min_cut_set = [], leak_impact_edges = []):
    G = H.copy()
    for u in ["super_source", "super_sink"]:
        if u in G.nodes():
            G.remove_node(u)
    nodes = []
    H = G.copy()
    cut_nodes = []
    for u, v in min_cut_set:
        if H[u][v]['capacity'] > 0:
            cut_nodes.append(u)    
    for u, v in min_cut_set:
        if H[u][v]['capacity'] > 0:
            H.remove_edge(u, v)
    nodes = [1,2,3]
    _, reachable_edges = get_outward_connected_components(H, nodes)
    edges = []
    for u, v in G.edges():
        edges.append((u,v))
    no_reachable_edges = set(edges) - reachable_edges
    """可视化管网及关阀方案"""
    plt.figure(figsize=(14, 10))
  
    # 按类型定义节点颜色
    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    H = G.copy()
    nodes = list(H.nodes)  # 创建静态的节点列表快照
    for node in nodes:
        attr = H.nodes[node]
        if attr['type'] == 'source':
            node_colors.append('lime')
            node_sizes.append(300)
            edgecolors.append("none")
            linewidths.append(0)     
        elif attr['type'] == 'leak':
            node_colors.append('red')
            node_sizes.append(300)  
            edgecolors.append("none")
            linewidths.append(0)                      
        elif attr['type'] == 'user':
            node_colors.append('gold')
            node_sizes.append(300)
            edgecolors.append("none")
            linewidths.append(0)            
        elif attr['type'] == 'other':
            node_colors.append('green')
            node_sizes.append(300)
            edgecolors.append("none")
            linewidths.append(0)            
        elif attr['type'] == 'valve' and node not in cut_nodes:
            node_colors.append('blue')
            node_sizes.append(300)
            edgecolors.append("gold")
            linewidths.append(2)            
        else:
            H.remove_node(node)
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9)

    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    H = G.copy()
    nodes = list(H.nodes)  # 创建静态的节点列表快照
    for node in nodes:
        attr = H.nodes[node]
        if attr['type'] == 'valve' and node in cut_nodes:
            node_colors.append('blue')
            node_sizes.append(300)
            edgecolors.append("gold")
            linewidths.append(2)
        else:
            H.remove_node(node)
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9, edgecolors = 'red', linewidths=3)
  
    # 绘制边
    edge_colors = []
    edge_widths = []
    for u, v in G.edges():
        if (u, v) not in no_reachable_edges:
            edge_colors.append('gray')
            edge_widths.append(1.5)
            
        else:
            edge_colors.append('red')
            edge_widths.append(1.5)
    nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, arrows=True, arrowstyle='-|>', arrowsize=20)

    if len(no_reachable_edges) > 0:
        ax = plt.gca()  # 获取当前坐标轴
        for u, v in no_reachable_edges:
            x1, y1 = pos[u]
            x2, y2 = pos[v]
            mid_x = (x1 + x2) / 2
            mid_y = (y1 + y2) / 2
            ax.plot(mid_x, mid_y, 'rx', markersize=15, markeredgewidth=2)  # 红色叉号
    if len(leak_impact_edges) > 0:
        print ("leak_impact_edges", leak_impact_edges)
        ax = plt.gca()  # 获取当前坐标轴            
        for u, v in leak_impact_edges:
            x1, y1 = pos[u]
            x2, y2 = pos[v]
            mid_x = (x1 + x2) / 2
            mid_y = (y1 + y2) / 2
            ax.plot(mid_x, mid_y, 'rx', markersize=15, markeredgewidth=2)  # 红色叉号

    # if len(min_cut_set) > 0:
    #     edge_colors = []
    #     edge_widths = []
    #     edge_styles = []
    #     for u, v in G.edges():
    #         if (u, v) in min_cut_set:
    #             edge_colors.append('red')
    #             edge_widths.append(1.5)
    #             edge_styles.append('dashed')
        
    #     nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, style=edge_styles, arrows=True, arrowstyle='-|>')
  
    # 节点标签
    labels = {}
    for node in G.nodes:
        attr = G.nodes[node]
        if attr['type'] == 'valve':
            labels[node] = f"V{node}"
        elif attr['type'] == 'user':
            labels[node] = f"U{node}"
        elif attr['type'] == 'leak':
            labels[node] = f"L{node}"
        elif attr['type'] == 'other':
            labels[node] = f"O{node}"            
        elif attr['type'] == 'source':
            labels[node] = f"S{node}"
        else:
            labels[node] = f"{node}"
  
    nx.draw_networkx_labels(G, pos, labels, font_size=9)
  
    # 管道属性标注
    for (u, v), attr in G.edges.items():
        plt.text((pos[u][0]+pos[v][0])/2, 
                 (pos[u][1]+pos[v][1])/2, 
                 f"{G[u][v].get('capacity')}", 
                 fontsize=7, color='blue')
  
    # 图例
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', label='气源(S)', markerfacecolor='lime', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='用户(U)', markerfacecolor='gold', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='泄漏点(L)', markerfacecolor='red', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='其他节点(O)', markerfacecolor='green', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)未关', markerfacecolor='blue', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)关闭', markerfacecolor='blue', markeredgecolor='red', markeredgewidth=3, markersize=10),
        #plt.Line2D([0], [0], marker='o', color='w', label='增加节点(超级源/超级汇)', markerfacecolor='blue', markersize=10),
        plt.Line2D([0], [0], color='gray', lw=1, label='普通管道'),
        plt.Line2D([0], [0], marker='x', color='red', linestyle='None', markersize=15, markeredgewidth=2, label='被阻断的管道')

    ]
  
    plt.legend(handles=legend_elements, loc='upper right', prop=font)
    plt.title(name, fontsize=14, fontproperties=font)
    plt.tight_layout()
    plt.savefig('%s.jpg' %name, dpi=300)
    plt.show()

def visualize(G, pos, valve_cut, users_impact):
    plt.figure(figsize=(14, 10))
    # 按类型定义节点颜色
    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    H = G.copy()
    for node in G.nodes():
        attr = G.nodes[node]
        if attr['type'] == 'source':
            node_colors.append('lime')
            node_sizes.append(300)
        elif attr['type'] == 'leak':
            node_colors.append('red')
            node_sizes.append(300)
        elif attr['type'] == 'user' and node not in users_impact:
            node_colors.append('gold')
            node_sizes.append(300)
        elif attr['type'] == 'other':
            node_colors.append('green')
            node_sizes.append(300)
        elif attr['type'] == 'valve' and node not in valve_cut:
            node_colors.append('blue')
            node_sizes.append(300)
        else:
            H.remove_node(node)
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9)

    node_colors = []
    node_sizes = []
    edgecolors = []
    linewidths = []
    H = G.copy()
    for node in G.nodes():
        attr = G.nodes[node]
        if attr['type'] == 'valve' and node in valve_cut:
            node_colors.append('blue')
            node_sizes.append(300)
            edgecolors.append("red")
            linewidths.append(2)
        elif attr['type'] == 'user' and node in users_impact:
            node_colors.append('gold')
            node_sizes.append(300)
            edgecolors.append("red")
            linewidths.append(2)            
        else:
            H.remove_node(node)
    # 绘制节点
    nx.draw_networkx_nodes(H, pos, node_size=node_sizes, node_color=node_colors, alpha=0.9, edgecolors = 'red', linewidths=3)
  
    # 绘制边
    edge_colors = []
    edge_widths = []
    for u, v in G.edges():
        edge_colors.append('gray')
        edge_widths.append(1.5)
    nx.draw_networkx_edges(G, pos, edge_color=edge_colors, width=edge_widths, arrows=True, arrowstyle='-|>', arrowsize=20)

    # 节点标签
    labels = {}
    for node in G.nodes:
        attr = G.nodes[node]
        if attr['type'] == 'valve':
            labels[node] = f"V{node}"
        elif attr['type'] == 'user':
            labels[node] = f"U{node}"
        elif attr['type'] == 'leak':
            labels[node] = f"L{node}"
        elif attr['type'] == 'other':
            labels[node] = f"O{node}"            
        elif attr['type'] == 'source':
            labels[node] = f"S{node}"
        else:
            labels[node] = f"{node}"
    nx.draw_networkx_labels(G, pos, labels, font_size=9)
  
    # 图例
    legend_elements = [
        plt.Line2D([0], [0], marker='o', color='w', label='气源(S)', markerfacecolor='lime', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='用户(U)不受影响', markerfacecolor='gold', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='用户(U)受影响', markerfacecolor='gold', markeredgecolor='red', markeredgewidth=3, markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='泄漏点(L)', markerfacecolor='red', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='其他节点(O)', markerfacecolor='green', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)未关', markerfacecolor='blue', markersize=10),
        plt.Line2D([0], [0], marker='o', color='w', label='阀门(V)关闭', markerfacecolor='blue', markeredgecolor='red', markeredgewidth=3, markersize=10),
        plt.Line2D([0], [0], color='gray', lw=1, label='普通管道'),
        plt.Line2D([0], [0], marker='x', color='red', linestyle='None', markersize=15, markeredgewidth=2, label='被阻断的管道')

    ]
  
    plt.legend(handles=legend_elements, loc='upper right', prop=font)
    #plt.title(name, fontsize=14, fontproperties=font)
    plt.tight_layout()
    #plt.savefig('%s.jpg' %name, dpi=300)
    plt.show()
    #plt.savefig("main_final.jpg", dpi=300)